"""
Regress one node on its parents with an Additive Noise Model
"""
import numpy as np
from sklearn.linear_model._base import LinearModel
from sklearn.linear_model import LinearRegression

from .base import Regressor


class ANMRegressor(Regressor):
    """
    加性噪声模型回归器 (Additive Noise Model Regressor)
    
    基于假设：x = f(pa(X)) + e，其中：
    - x: 当前节点的观测值
    - pa(X): 父节点的观测值集合
    - f: 父节点到当前节点的映射函数（这里假设为线性函数）
    - e: 服从正态分布的噪声项
    """

    def __init__(self, regressor: LinearModel = None, **kwargs):
        """
        初始化回归器
        
        Args:
            regressor: sklearn线性模型实例，默认使用普通线性回归
            **kwargs: 传递给父类的额外参数
        """
        super().__init__(**kwargs)
        # 如果没有指定回归器，则使用默认的线性回归
        self._regressor = regressor if regressor else LinearRegression()

    def _score(
        self,
        train_x: np.ndarray,    # 训练集的父节点特征矩阵
        test_x: np.ndarray,     # 测试集的父节点特征矩阵
        train_y: np.ndarray,    # 训练集的当前节点值
        test_y: np.ndarray,     # 测试集的当前节点值
    ) -> np.ndarray:
        """
        计算异常分数
        
        工作流程：
        1. 使用训练数据拟合线性模型，学习f(pa(X))
        2. 计算训练集和测试集的预测误差（即噪声项e）
        3. 使用Z分数标准化误差，得到异常分数
        
        Returns:
            np.ndarray: 标准化后的异常分数，分数越高表示异常程度越大
        """
        # 使用训练数据拟合模型
        self._regressor.fit(train_x, train_y)
        # 计算训练集的预测误差（残差）
        train_err: np.ndarray = train_y - self._regressor.predict(train_x)
        # 计算测试集的预测误差（残差）
        test_err: np.ndarray = test_y - self._regressor.predict(test_x)
        # 将误差转换为标准化的Z分数
        return self._zscore(train_y=train_err, test_y=test_err)
